import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
import cv2
import numpy as np
from datetime import datetime
import os

class ImageSaver(Node):
    def __init__(self, save_dir='./out/saved_images'):
        super().__init__('jpeg_saver')
        self.subscription = self.create_subscription(
            CompressedImage,
            '/camera/image/compressed',
            self.image_callback,
            10)

        # 创建保存目录（如果不存在）
        self.save_dir = save_dir
        os.makedirs(self.save_dir, exist_ok=True)
        self.get_logger().info(f'Images will be saved to: {os.path.abspath(self.save_dir)}')

        # 初始化计数器（用于生成唯一文件名）
        self.image_counter = 0

    def image_callback(self, msg):
        try:
            # 解码JPEG数据
            np_arr = np.frombuffer(msg.data, np.uint8)
            img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

            # 生成唯一文件名（两种方式任选）
            # 方式1：时间戳 + 计数器
            timestamp = datetime.now().strftime("%Y%m%d-%H%M%S%f")
            filename = f"{self.save_dir}/image_{timestamp}_{self.image_counter}.jpg"

            # 方式2：纯计数器（适合固定频率场景）
            # filename = f"{self.save_dir}/image_{self.image_counter:06d}.jpg"

            # 保存图像
            cv2.imwrite(filename, img)
            self.get_logger().info(f'Saved: {filename} (Size: {img.shape[1]}x{img.shape[0]})')

            # 更新计数器
            self.image_counter += 1

        except Exception as e:
            self.get_logger().error(f'Error saving image: {str(e)}')

def main(args=None):
    rclpy.init(args=args)
    node = ImageSaver()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()